#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
直接下载模型权重
"""

import os
import sys
import requests
import argparse
from tqdm import tqdm

# 改为提供百度网盘链接，用户需要手动下载
DOWNLOAD_INSTRUCTIONS = {
    'u2net': {
        'url': 'https://github.com/xuebinqin/U-2-Net#usage-for-salient-object-detection',
        'baidu_url': 'https://pan.baidu.com/s/1-R_BIPSmGXEMZX0j4OeLng 提取码: pf9k',
        'size': '176.3 MB'
    },
    'u2netp': {
        'url': 'https://github.com/xuebinqin/U-2-Net#usage-for-salient-object-detection',
        'baidu_url': 'https://pan.baidu.com/s/1PHfBwo1tJUEQCUDlmPVQaQ 提取码: 8xsi',
        'size': '4.7 MB'
    }
}

def provide_download_instructions(model_name):
    """
    提供下载说明
    
    Args:
        model_name (str): 模型名称 ('u2net', 'u2netp')
    """
    if model_name not in DOWNLOAD_INSTRUCTIONS:
        print(f"不支持的模型: {model_name}")
        print(f"支持的模型: {', '.join(DOWNLOAD_INSTRUCTIONS.keys())}")
        return False
    
    # 确定模型文件路径
    if model_name == 'u2net':
        model_filename = 'u2net_model.pth'
    elif model_name == 'u2netp':
        model_filename = 'u2netp_model.pth'
    else:
        raise ValueError(f"不支持的模型名称: {model_name}")
    
    # 获取模型文件夹路径
    model_dir = os.path.join(os.getcwd(), 'models', 'weights')
    # 确保目录存在
    os.makedirs(model_dir, exist_ok=True)
    
    # 完整的模型文件路径
    model_path = os.path.join(model_dir, model_filename)
    
    # 检查模型是否已经存在
    if os.path.exists(model_path):
        print(f"模型已存在: {model_path}")
        return True
    
    info = DOWNLOAD_INSTRUCTIONS[model_name]
    print(f"\n=============== {model_name} 模型下载说明 ===============")
    print(f"模型大小: {info['size']}")
    print(f"GitHub仓库地址: {info['url']}")
    print(f"百度网盘链接: {info['baidu_url']}")
    print(f"\n请手动下载模型文件并放置到以下路径:")
    print(f"{model_path}")
    print("=================================================\n")
    
    return True

def main():
    parser = argparse.ArgumentParser(description="模型权重下载说明工具")
    parser.add_argument("model", nargs="?", default="all", 
                        help="要下载的模型名称 (u2net, u2netp, all)")
    
    args = parser.parse_args()
    model_name = args.model.lower()
    
    if model_name == "all":
        print("显示所有模型的下载说明...")
        for model in DOWNLOAD_INSTRUCTIONS.keys():
            provide_download_instructions(model)
    elif model_name in DOWNLOAD_INSTRUCTIONS.keys():
        provide_download_instructions(model_name)
    else:
        print(f"不支持的模型: {model_name}")
        print(f"支持的模型: {', '.join(DOWNLOAD_INSTRUCTIONS.keys())}, all")
        return 1
    
    return 0

if __name__ == "__main__":
    main() 